///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file command.cpp
///
/// \brief   file with implementation for functions and classes related to commands
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<src/pcps.h>
#include<src/command.h>
#include<src/all_data.h>
#include<src/mpi_interface.h>
#include<src/correlator.h>
#include<src/corr_prod.h>
#include<src/restricted_pairing_function.h>
#include<src/hamiltonian.h>
#include<src/hf.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  The execution of a comment command ignores the rest of the text on an input file line.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::CommentCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  // get MPI info
  const int nproc  = pcps::mpi::size();
  const int myrank = pcps::mpi::rank();

  // read in and discard the rest of the line on the root node
  std::string comment;
  if (myrank == 0)
    arch.getline(comment);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  The execution of an OptionCommand sets one of the user input options.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::OptionCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  // get MPI info
  const int nproc  = pcps::mpi::size();
  const int myrank = pcps::mpi::rank();

  // get the member name for the user option
  std::string member_name;
  pcps::mpi::read_and_bcast(arch, member_name, "Failed to read member name when executing OptionCommand.");

  // set the member's value
  data.userinp.set_member(member_name, arch);

  // verify user options
  data.userinp.verify();

  // re-initialize derived members
  data.userinp.initialize_derived_data();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  The execution of a WfnDataCommand adds a WfnData object of the appropriate
///         type to either the real or complex wavefunction.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class T> void pcps::WfnDataCommand<T>::execute(pcps::Archive & arch, pcps::AllData & data) {

  // integer
  if (typeid(T) == typeid(int)) {

    if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real)
      data.wd.add_wfn_int_data(arch);
    else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex)
      data.wc.add_wfn_int_data(arch);
    else 
      throw pcps::Exception("Unknown scalar type \"%s\" in pcps::WfnDataCommand::execute") % data.userinp.scalar_type();

  // real
  } else if (typeid(T) == typeid(double)) {

    data.wd.add_wfn_float_data(arch);

  // complex
  } else if (typeid(T) == typeid(std::complex<double>)) {

    data.wc.add_wfn_float_data(arch);

  // other
  } else 
    throw pcps::Exception("Unknown T type \"%s\" in pcps::WfnDataCommand::execute") % typeid(T).name();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  The execution of a OverallCoeffCommand sets the overall wavefunction coefficient.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::OverallCoeffCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  // get the name of the overall coefficient to use
  std::string s;
  pcps::mpi::read_and_bcast(arch, s, "Failed to read coefficient name in pcps::OverallCoeffCommand::execute.");

  // set the overall coefficient
  if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real)
    data.wd.set_overall_coefficient(s);
  else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex)
    data.wc.set_overall_coefficient(s);
  else 
    throw pcps::Exception("Unknown scalar type \"%s\" in pcps::OverallCoeffCommand::execute") % data.userinp.scalar_type();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  The execution of a CorrelatorCommand adds a Correlator object to the wavefunction.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::CorrelatorCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real)
    data.wd.add_coeff_obj<pcps::Correlator<double> >(data.userinp, arch);
  else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex)
    data.wc.add_coeff_obj<pcps::Correlator<std::complex<double> > >(data.userinp, arch);
  else 
    throw pcps::Exception("Unknown scalar type \"%s\" in pcps::CorrelatorCommand::execute") % data.userinp.scalar_type();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  The execution of a CorrProdCommand adds a CorrProd object to the wavefunction.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::CorrProdCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real)
    data.wd.add_coeff_obj<pcps::CorrProd<double> >(data.userinp, arch);
  else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex)
    data.wc.add_coeff_obj<pcps::CorrProd<std::complex<double> > >(data.userinp, arch);
  else 
    throw pcps::Exception("Unknown scalar type \"%s\" in pcps::CorrProdCommand::execute") % data.userinp.scalar_type();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  The execution of a RestPairingFuncCommand adds a RestPairingFunc object to the wavefunction.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::RestPairingFuncCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real)
    data.wd.add_coeff_obj<pcps::RestPairingFunc<double> >(data.userinp, arch);
  else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex)
    data.wc.add_coeff_obj<pcps::RestPairingFunc<std::complex<double> > >(data.userinp, arch);
  else 
    throw pcps::Exception("Unknown scalar type \"%s\" in pcps::RestPairingFuncCommand::execute") % data.userinp.scalar_type();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  The execution of a LinearMethodCommand optimizes the wavefunction with the fast
///         linear method.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::LinearMethodCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real)
    pcps::linear_method_optimization(data.userinp, data.hd, data.wd);
  else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex)
    pcps::linear_method_optimization(data.userinp, data.hc, data.wc);
  else 
    throw pcps::Exception("Unknown scalar type \"%s\" in pcps::LinearMethodCommand::execute") % data.userinp.scalar_type();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Loads the one and two electron integrals from psi3.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::Psi3IntsCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real)
    data.hd.initialize_integrals(data.userinp, "psi3");
  else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex)
    data.hc.initialize_integrals(data.userinp, "psi3");
  else 
    throw pcps::Exception("Unknown scalar type \"%s\" in pcps::Psi3IntsCommand::execute") % data.userinp.scalar_type();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Loads the one and two electron integrals from molpro.
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::MolproIntsCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real)
    data.hd.initialize_integrals(data.userinp, "molpro");
  else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex)
    data.hc.initialize_integrals(data.userinp, "molpro");
  else 
    throw pcps::Exception("Unknown scalar type \"%s\" in pcps::MolproIntsCommand::execute") % data.userinp.scalar_type();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Runs a restricted hartree fock calculation
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::RHFCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  // get the name of the wavefunction data to use
  std::string name;
  pcps::mpi::read_and_bcast(arch, name, "Failed to read data name for restricted hartree fock orbital coefficients.");

  // ensure the number of particles is even
  if ( data.userinp.nparticles() % 2 != 0 )
    throw pcps::Exception("restricted hartree fock requires an even number of particles");

  // run the hartree fock calculation
  if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real) {

    // ensure the data vector is the correct size
    if ( data.wd.get_fdata(name).data_vec().size() != data.userinp.n_act_orb() * data.userinp.n_act_orb() )
      throw pcps::Exception("data vector size incorrect in pcps::RHFCommand::execute.");

    double energy;
    pcps::hf::run_scf(true,
                      data.userinp.n_act_orb(),
                      data.userinp.nparticles()/2,
                      data.userinp.nparticles()/2,
                      data.hd,
                      data.userinp.uhf_bias(),
                      data.userinp.uhf_bias_decay(),
                      data.userinp.hf_thresh(),
                      data.userinp.max_hf_iter(),
                      data.wd.get_fdata(name).data_vec(),
                      data.wd.get_fdata(name).data_vec(),
                      energy);

  } else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex) {

    // ensure the data vector is the correct size
    if ( data.wc.get_fdata(name).data_vec().size() != data.userinp.n_act_orb() * data.userinp.n_act_orb() )
      throw pcps::Exception("data vector size incorrect in pcps::RHFCommand::execute.");

    std::complex<double> energy;
    pcps::hf::run_scf(true,
                      data.userinp.n_act_orb(),
                      data.userinp.nparticles()/2,
                      data.userinp.nparticles()/2,
                      data.hc,
                      data.userinp.uhf_bias(),
                      data.userinp.uhf_bias_decay(),
                      data.userinp.hf_thresh(),
                      data.userinp.max_hf_iter(),
                      data.wc.get_fdata(name).data_vec(),
                      data.wc.get_fdata(name).data_vec(),
                      energy);

  } else 
    throw pcps::Exception("Unknown scalar type \"%s\" in pcps::RHFCommand::execute") % data.userinp.scalar_type();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Runs an unrestricted hartree fock calculation
///
/// \param[in,out] arch      The archive from which the command execution may read instructions.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void pcps::UHFCommand::execute(pcps::Archive & arch, pcps::AllData & data) {

  // get the name of the wavefunction data to use for the alpha coefficients
  std::string a_name;
  pcps::mpi::read_and_bcast(arch, a_name, "Failed to read data name for unrestricted hartree fock alpha orbital coefficients.");

  // get the name of the wavefunction data to use for the beta coefficients
  std::string b_name;
  pcps::mpi::read_and_bcast(arch, b_name, "Failed to read data name for unrestricted hartree fock beta orbital coefficients.");

  // run the hartree fock calculation
  if (data.userinp.scalar_type() == pcps::option_value::Scalar_Real) {

    // ensure the data vectors are the correct size
    if ( data.wd.get_fdata(a_name).data_vec().size() != data.userinp.n_act_orb() * data.userinp.n_act_orb() )
      throw pcps::Exception("alpha coefficient matrix size incorrect in pcps::UHFCommand::execute.");
    if ( data.wd.get_fdata(b_name).data_vec().size() != data.userinp.n_act_orb() * data.userinp.n_act_orb() )
      throw pcps::Exception("beta coefficient matrix size incorrect in pcps::UHFCommand::execute.");

    double energy;
    pcps::hf::run_scf(false,
                      data.userinp.n_act_orb(),
                      pcps::round( 0.5 * data.userinp.nparticles() + data.userinp.z_spin() ),
                      pcps::round( 0.5 * data.userinp.nparticles() - data.userinp.z_spin() ),
                      data.hd,
                      data.userinp.uhf_bias(),
                      data.userinp.uhf_bias_decay(),
                      data.userinp.hf_thresh(),
                      data.userinp.max_hf_iter(),
                      data.wd.get_fdata(a_name).data_vec(),
                      data.wd.get_fdata(b_name).data_vec(),
                      energy);

  } else if (data.userinp.scalar_type() == pcps::option_value::Scalar_Complex) {

    // ensure the data vectors are the correct size
    if ( data.wc.get_fdata(a_name).data_vec().size() != data.userinp.n_act_orb() * data.userinp.n_act_orb() )
      throw pcps::Exception("alpha coefficient matrix size incorrect in pcps::UHFCommand::execute.");
    if ( data.wc.get_fdata(b_name).data_vec().size() != data.userinp.n_act_orb() * data.userinp.n_act_orb() )
      throw pcps::Exception("beta coefficient matrix size incorrect in pcps::UHFCommand::execute.");


    std::complex<double> energy;
    pcps::hf::run_scf(true,
                      data.userinp.n_act_orb(),
                      pcps::round( 0.5 * data.userinp.nparticles() + data.userinp.z_spin() ),
                      pcps::round( 0.5 * data.userinp.nparticles() - data.userinp.z_spin() ),
                      data.hc,
                      data.userinp.uhf_bias(),
                      data.userinp.uhf_bias_decay(),
                      data.userinp.hf_thresh(),
                      data.userinp.max_hf_iter(),
                      data.wc.get_fdata(a_name).data_vec(),
                      data.wc.get_fdata(b_name).data_vec(),
                      energy);

  } else 
    throw pcps::Exception("Unknown scalar type \"%s\" in pcps::UHFCommand::execute") % data.userinp.scalar_type();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Returns a shared pointer pointing to a newly created command of the type whose
///         keyword matches the supplied key.
///
/// \param[in]     key       The keyword of the desired command type.
///
/// \return An empty (default constructed) command of the requested type.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
boost::shared_ptr<pcps::Command> pcps::get_command(const std::string & key) {

  // create a short name for the return type
  typedef boost::shared_ptr<pcps::Command> ptr;

  // initialize the return value
  ptr retval;

  // point to an object of the appropriate derived type
  if      ( key == pcps::BlankCommand::keyword()                          ) retval = ptr( new pcps::BlankCommand()                          );
  else if ( key == pcps::CommentCommand::keyword()                        ) retval = ptr( new pcps::CommentCommand()                        );
  else if ( key == pcps::OptionCommand::keyword()                         ) retval = ptr( new pcps::OptionCommand()                         );
  else if ( key == pcps::WfnDataCommand<int>::keyword()                   ) retval = ptr( new pcps::WfnDataCommand<int>()                   );
  else if ( key == pcps::WfnDataCommand<double>::keyword()                ) retval = ptr( new pcps::WfnDataCommand<double>()                );
  else if ( key == pcps::WfnDataCommand<std::complex<double> >::keyword() ) retval = ptr( new pcps::WfnDataCommand<std::complex<double> >() );
  else if ( key == pcps::OverallCoeffCommand::keyword()                   ) retval = ptr( new pcps::OverallCoeffCommand()                   );
  else if ( key == pcps::CorrelatorCommand::keyword()                     ) retval = ptr( new pcps::CorrelatorCommand()                     );
  else if ( key == pcps::CorrProdCommand::keyword()                       ) retval = ptr( new pcps::CorrProdCommand()                       );
  else if ( key == pcps::RestPairingFuncCommand::keyword()                ) retval = ptr( new pcps::RestPairingFuncCommand()                );
  else if ( key == pcps::LinearMethodCommand::keyword()                   ) retval = ptr( new pcps::LinearMethodCommand()                   );
  else if ( key == pcps::Psi3IntsCommand::keyword()                       ) retval = ptr( new pcps::Psi3IntsCommand()                       );
  else if ( key == pcps::MolproIntsCommand::keyword()                     ) retval = ptr( new pcps::MolproIntsCommand()                     );
  else if ( key == pcps::RHFCommand::keyword()                            ) retval = ptr( new pcps::RHFCommand()                            );
  else if ( key == pcps::UHFCommand::keyword()                            ) retval = ptr( new pcps::UHFCommand()                            );
  else
    throw pcps::Exception("Unknown command keyword \"%s\"") % key;

  // return the empty command of the given type
  return retval;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Execute all commands contained within an archive file.
///         The template parameter T is the archive type, which must be either
///         pcps::TextFileArchive or pcps::BinaryFileArchive.
///
/// \param[in]     filename  The name of the file.
/// \param[in,out] data      An object holding the data on which the command will act.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class T> void pcps::run_commands(const std::string & filename, pcps::AllData & data) {

  // get MPI info
  const int nproc  = pcps::mpi::size();
  const int myrank = pcps::mpi::rank();

  // open the file on the root process
  boost::shared_ptr<pcps::Archive> f;
  if (myrank == 0)
    f = boost::shared_ptr<pcps::Archive>( new T(filename, "r") );
  else
    f = boost::shared_ptr<pcps::Archive>( new pcps::BinaryArchive() );

  // one by one, execute the commands in the file
  while (true) {

    // read the keyword
    std::string keyword;
    if (myrank == 0)
      *f >> keyword;
    pcps::mpi::bcast(keyword);

    // if there were no more keywords, we are finished
    bool finished = !*f;
    pcps::mpi::bcast(finished);
    if (finished) break;

    // run the command corresponding to the keyword
    pcps::get_command(keyword)->execute(*f, data);

  }

}

// instantiate the template
template void pcps::run_commands<pcps::TextFileArchive>(const std::string &, pcps::AllData &);
template void pcps::run_commands<pcps::BinaryFileArchive>(const std::string &, pcps::AllData &);
